function run_modQTL_imputation_spikeslab(YORG,J,sigsq_prior,B0,INITTIME,BURNIN,NGIBBS,BATCHSIZE,OUTDIR)
%
% run_modQTL_imputation_spikeslab(YORG,J,SIGSQ_PRIOR,B0,INITTIME,BURNIN,NGIBBS,BATCHSIZE,OUTDIR)
%
% YORG       = ntis x nind x ngene (incomplete) expression matrix
% J          = number of factors
% SIGSQ_PRIOR= prior for sigma_observed * sigsq_prior (scaling)
% B0         = prior for beta(1,b0)
% BURNIN     = burn-in iterations wihtout sparsity
% NGIBBS     = sampling time
% BATCHSIZE  = size of minibatch
% OUTDIR     = directory in which result files are stored
%
% code: Yongjin Park, ypp@csail.mit.edu
%

    ok = mkdir(OUTDIR);

    if ~ok,
        fprintf(2,'Cannot create the directory: %s\n', OUTDIR);
        exit;
    end

    if OUTDIR(end) == '/',
        OUTDIR = OUTDIR(1:(end-1));
    end

    [Ntis,Nind,Ngene] = size(YORG);

    A0 = 1;

    prior_odds_tis = zeros(Ntis,1,'single');
    prior_odds_factor = 0;


    % ****************************************************************
    % initial parameters
    SIGSQ0 = var(YORG(~isnan(YORG)));
    sigsq = 0.1*SIGSQ0;
    U = 0.1*randn(Ntis,J,'single');
    invtausq_u = 0.01*Nind;


    % Bayesian model averaging
    u_spike_samples = NaN(Ntis,J,NGIBBS,'single');
    u_slab_samples = NaN(Ntis,J,NGIBBS,'single');

    Vavg = zeros(J,Nind,Ngene,'single');
    VZavg = zeros(J,Nind,Ngene,'single');
    rank_spike_avg = zeros(Ngene,J,'single');
    % rank_spike_samples = NaN(Ngene,J,NGIBBS,'single');
    Z_factor = zeros(J,Ngene,'single');
    Z_tis = ones(Ntis,J,'single');
    slab_tis = U;

    intercept_avg = zeros(Ntis,Ngene,'single');
    % invtausq_v_samples = NaN(Ngene,J,NGIBBS,'single');

    % 1. construct minibatches
    Nbatch = ceil(Ngene / BATCHSIZE);
    permuted_genes = randsample(Ngene,Ngene);

    BURNIN = BURNIN * Nbatch;
    NGIBBS = NGIBBS * Nbatch;
    INITTIME = INITTIME * Nbatch;

    GENES_minibatch = arrayfun(@(b) permuted_genes((1+b*BATCHSIZE):min(Ngene,(b+1)*BATCHSIZE)), ...
                               (0:(Nbatch-1))',...
                               'UniformOutput',false);

    T_minibatch = cell(Nbatch,1); T_minibatch(:) = {zeros(Ntis,J,'single')};
    S_minibatch = cell(Nbatch,1); S_minibatch(:) = {zeros(J,J,'single')};
    RSS_minibatch = zeros(Nbatch,1,'single');
    OBS_minibatch = zeros(Nbatch,1,'single');


    missing_minibatch = cellfun(@(genes) reshape(isnan(YORG(:,:,genes)),Ntis,Nind*numel(genes)), ...
                                GENES_minibatch, 'UniformOutput', false);

    Yimpute_minibatch = cellfun(@(genes) reshape(YORG(:,:,genes),Ntis,Nind*numel(genes)),...
                                GENES_minibatch, 'UniformOutput', false);

    for bb = 0:(Nbatch-1),
        % remove NaN
        Yimpute_minibatch{bb+1}( missing_minibatch{bb+1} ) = 0;
    end    

    Vt_minibatch = cellfun(@(genes) zeros(Nind*numel(genes),J,'single'), ...
                           GENES_minibatch, 'UniformOutput', false);

    % Zt_minibatch = cellfun(@(genes) ones(Nind*numel(genes),J,'single'), ...
    %                       GENES_minibatch, 'UniformOutput', false);

    invtausq_v_minibatch = cellfun(@(genes) 0.01*Ntis*ones(numel(genes),J,'single'), ...
                                   GENES_minibatch, 'UniformOutput', false);

    intercept_minibatch = cellfun(@(genes) zeros(Ntis,numel(genes),'single'), GENES_minibatch, 'UniformOutput', false);

    % ****************************************************************
    % 1 x gene <-> 1 x gene*ind    
    gene2geneind = single(full(sparse(reshape(repmat(1:BATCHSIZE,Nind,1),1,[]),1:(BATCHSIZE*Nind),1)));

    % ****************************************************************
    % locally estimate V and
    % make the global statistics exact
    GLOBAL_T = zeros(Ntis,J,'single'); % sum Y * V'
    GLOBAL_S = zeros(J,J,'single');    % sum V * V'

    GLOBAL_RSS = 0;
    GLOBAL_OBS = 0;

    UtU = U'*U;

    for bb = 0:(Nbatch-1),

        Y = Yimpute_minibatch{bb+1}; % tis x (ind * gene)
        Vt = Vt_minibatch{bb+1};     % (ind * gene) x J

        % Zt = Zt_minibatch{bb+1};     % (ind * gene) x J

        intercept = intercept_minibatch{bb+1};  % tis x gene
        invtausq_v = invtausq_v_minibatch{bb+1}; % gene x J

        genes = permuted_genes( (1+bb*BATCHSIZE):min(Ngene,(bb+1)*BATCHSIZE) );
        
        % gene*ind -> gene for this batch
        geneind2gene_b = gene2geneind(1:numel(genes), 1:size(Vt,1))';

        % update intercept
        intercept = (Y - U*Vt') * geneind2gene_b / (Nind); % tis x gene

        % precomputation
        YtU = (Y - intercept * geneind2gene_b')'*U;
        % YtU = Y'*U;
        Zt = Z_factor(:,genes)'; % gene x J
        VZtUtU = (Vt .* (geneind2gene_b * Zt)) * UtU; % (ind*gene) x J
        
        % VZtUtU = (Vt .* Zt) * UtU; % (ind*gene) x J
        Vsq = geneind2gene_b' * (Vt.^2); % gene x J

        tau_repeat = ceil([1:size(invtausq_v,1)*Nind]/Nind);

        for k = 1:J,
            vt_old = Vt(:,k);
            zz_old = Zt(:,k);

            % **************** sample slab ****************
            denom_v = UtU(k,k) + invtausq_v(tau_repeat,k);
            num_v = YtU(:,k) - VZtUtU(:,k) + (geneind2gene_b*zz_old) .* Vt(:,k) * UtU(k,k);
            % num_v = YtU(:,k) - VZtUtU(:,k) + zz_old .* Vt(:,k) * UtU(k,k);
            vt = num_v ./ denom_v + randn(numel(num_v),1,'single').*sqrt(sigsq./denom_v);

            % update precomputation
            delta_vz = (vt - vt_old).*(geneind2gene_b*zz_old);
            % delta_vz = (vt - vt_old).*zz_old;
            VZtUtU = VZtUtU + bsxfun(@times, delta_vz, UtU(:,k).');
            Vsq(:,k) = geneind2gene_b' * (vt.^2);

            Vt(:,k) = vt;

            % **************** sample spike ****************
            mass = geneind2gene_b' * (YtU(:,k).*Vt(:,k));
            mass = mass - geneind2gene_b' * (VZtUtU(:,k).*Vt(:,k));
            mass = mass + UtU(k,k) * (zz_old - 0.5) .* Vsq(:,k);

            % mass = YtU(:,k).*Vt(:,k);
            % mass = mass - VZtUtU(:,k).*Vt(:,k);
            % mass = mass + UtU(k,k) * (zz_old - 0.5) .* (vt.^2);

            zz = 1./(1+exp(-mass/sigsq));
            zz = single(rand(size(zz),'single') < zz);

            % update precomputation
            Zt(:,k) = zz;
            delta_vz = Vt(:,k).*(geneind2gene_b * (zz - zz_old));
            % delta_vz = Vt(:,k).*(zz - zz_old);

            VZtUtU = VZtUtU + bsxfun(@times, delta_vz, UtU(:,k).');
        end

        intercept_minibatch{bb+1} = intercept;

        VZt = Vt .* (geneind2gene_b * Zt);
        % VZt = Vt .* Zt;
        T_minibatch{bb+1} = (Y - intercept * geneind2gene_b')* VZt; % tis x J

        %- T_minibatch{bb+1} = Y * Vt; % tis x J
        S_minibatch{bb+1} = Vt' * Vt; % J x J
        Vt_minibatch{bb+1} = Vt;
        % Zt_minibatch{bb+1} = Zt;

        GLOBAL_T = GLOBAL_T + T_minibatch{bb+1};
        GLOBAL_S = GLOBAL_S + S_minibatch{bb+1};

        missing = missing_minibatch{bb+1};
        resid = Y - U*Vt';

        RSS_minibatch(bb+1) = sum(resid(~missing).^2);
        OBS_minibatch(bb+1) = sum(~missing(:));
        GLOBAL_RSS = GLOBAL_RSS + RSS_minibatch(bb+1);
        GLOBAL_OBS = GLOBAL_OBS + OBS_minibatch(bb+1);

        % update global factor usage
        Z_factor(:,genes) = Zt';
        % Z_factor(:,genes) = single(Zt' * geneind2gene_b > 0.5 * Nind);

        fprintf('Constructing minibatch b = %03d\r', bb);
    end

    % refit U matrix
    GLOBAL_N = sum(Z_factor,2);
    for jj=1:J
        uu_old = U(:,jj);
        tt = GLOBAL_T(:,jj);
        ss = GLOBAL_S(:,jj);
        
        denom_u = ss(jj) + invtausq_u*GLOBAL_N(jj);

        if denom_u <= 1e-4, % avoid NaN
            denom_u = 1e-4;
        end

        uu = (tt - U*ss + uu_old*ss(jj))/denom_u;
        uu = uu + randn(Ntis,1,'single') .* sqrt(sigsq/denom_u);
        
        U(:,jj) = uu;
    end


    % ****************************************************************
    % 2. online update
    rmse_trace = NaN(BURNIN+NGIBBS+INITTIME,1,'single');
    for iter = 1:(BURNIN + NGIBBS + INITTIME),

        bb = mod(iter,Nbatch);
        epoch = floor(iter/Nbatch);
        gibbs = ceil((iter - BURNIN - INITTIME)/Nbatch);

        % local inference
        Y = Yimpute_minibatch{bb+1};
        Vt = Vt_minibatch{bb+1};
        % Zt = Zt_minibatch{bb+1};     % (ind * gene) x J

        intercept = intercept_minibatch{bb+1};  % tis x gene
        invtausq_v = invtausq_v_minibatch{bb+1}; % gene x J

        genes = permuted_genes( (1+bb*BATCHSIZE):min(Ngene,(bb+1)*BATCHSIZE) );
        ng = numel(genes);
        geneind2gene_b = gene2geneind(1:ng, 1:size(Vt,1))';

        % update intercept
        intercept = (Y - U*Vt') * geneind2gene_b / (Nind); % tis x gene

        % subtract out local effects from the global
        GLOBAL_T = GLOBAL_T - T_minibatch{bb+1};
        GLOBAL_S = GLOBAL_S - S_minibatch{bb+1};
        GLOBAL_RSS = GLOBAL_RSS - RSS_minibatch(bb+1);
        GLOBAL_OBS = GLOBAL_OBS - OBS_minibatch(bb+1);

        % precomputation
        UtU = U'*U;
        YtU = (Y - intercept * geneind2gene_b')'*U;
        % YtU = Y'*U;
        Zt = Z_factor(:,genes)'; % gene x J
        VZtUtU = (Vt .* (geneind2gene_b * Zt)) * UtU; % (ind*gene) x J

        r = betarnd(A0+sum(Zt,2),B0+sum(1-Zt,2));
        prior_odds_factor = log(r) - log(1-r);

        % VZtUtU = (Vt .* Zt) * UtU; % (ind*gene) x J
        Vsq = geneind2gene_b' * (Vt.^2); % gene x J
        VsqDeterm = zeros(size(Vsq),'single');

        tau_repeat = ceil([1:size(invtausq_v,1)*Nind]/Nind);

        for k = 1:J,
            vt_old = Vt(:,k);
            zz_old = Zt(:,k);

            % **************** sample slab ****************
            denom_v = UtU(k,k) + invtausq_v(tau_repeat,k);
            num_v = YtU(:,k) - VZtUtU(:,k) + (geneind2gene_b*zz_old) .* Vt(:,k) * UtU(k,k);
            % num_v = YtU(:,k) - VZtUtU(:,k) + zz_old .* Vt(:,k) * UtU(k,k);
            vt = num_v ./ denom_v;
            VsqDeterm(:,k) = geneind2gene_b' * (vt.^2);
            vt = vt + randn(numel(num_v),1,'single').*sqrt(sigsq./denom_v);

            % update precomputation
            delta_vz = (vt - vt_old).*(geneind2gene_b*zz_old);
            % delta_vz = (vt - vt_old).* zz_old;
            VZtUtU = VZtUtU + bsxfun(@times, delta_vz, UtU(:,k).');
            Vsq(:,k) = geneind2gene_b' * (vt.^2);

            Vt(:,k) = vt;

            % debug = (Vt .* (geneind2gene_b * Zt)) * UtU; % (ind*gene) x J
            % assert( all(abs(debug(:) - VZtUtU(:))<1e-4) );

            % **************** sample spike ****************
            mass = geneind2gene_b' * (YtU(:,k).*Vt(:,k));
            mass = mass - geneind2gene_b' * (VZtUtU(:,k).*Vt(:,k));
            mass = mass + UtU(k,k) * (zz_old - 0.5) .* Vsq(:,k);

            % mass = YtU(:,k).*Vt(:,k);
            % mass = mass - VZtUtU(:,k).*Vt(:,k);
            % mass = mass + UtU(k,k) * (zz_old - 0.5) .* (vt.^2);

            zz = 1./(1+exp(- mass/sigsq - prior_odds_factor));
            zz = single(rand(size(zz),'single') < zz);

            % update precomputation
            Zt(:,k) = zz;
            delta_vz = Vt(:,k).*(geneind2gene_b * (zz - zz_old));
            % delta_vz = Vt(:,k).*(zz - zz_old);

            VZtUtU = VZtUtU + bsxfun(@times, delta_vz, UtU(:,k).');

            % debug = (Vt .* (geneind2gene_b * Zt)) * UtU; % (ind*gene) x J
            % assert( all(abs(debug(:) - VZtUtU(:))<1e-4) );
        end

        VZt = Vt .* (geneind2gene_b * Zt);
        % VZt = Vt .* Zt;
        intercept_minibatch{bb+1} = intercept;

        % locally estimate regularization
        if iter > INITTIME,
            invtausq_v = arrayfun(@(x) max(0.1/Nind,min(Nind*10,sigsq*Nind/x)), VsqDeterm);
        end

        invtausq_v_minibatch{bb+1} = invtausq_v;

        % imputation
        Yhat = U * VZt' + intercept*geneind2gene_b';
        % Yhat = U * VZt';
        missing = missing_minibatch{bb+1};
        Y(missing) = Yhat(missing);

        T_minibatch{bb+1} = (Y - intercept * geneind2gene_b')* VZt;
        % T_minibatch{bb+1} = Y * VZt; % tis x J
        S_minibatch{bb+1} = VZt' * VZt; % J x J

        Yimpute_minibatch{bb+1} = Y;
        Vt_minibatch{bb+1} = Vt;

        GLOBAL_T = GLOBAL_T + T_minibatch{bb+1};
        GLOBAL_S = GLOBAL_S + S_minibatch{bb+1};

        missing = missing_minibatch{bb+1};
        resid = Y - U*VZt' - intercept*geneind2gene_b';
        % resid = Y - U*VZt';

        RSS_minibatch(bb+1) = sum(resid(~missing).^2);
        OBS_minibatch(bb+1) = sum(~missing(:));
        GLOBAL_RSS = GLOBAL_RSS + RSS_minibatch(bb+1);
        GLOBAL_OBS = GLOBAL_OBS + OBS_minibatch(bb+1);

        Z_factor(:,genes) = Zt';
        % Z_factor(:,genes) = single(Zt' * geneind2gene_b > 0.5 * Nind);

        rmse = sqrt(GLOBAL_RSS / GLOBAL_OBS);
        rmse_trace(iter) = rmse;


        % ****************************************************************
        % refit U matrix
        GLOBAL_N = sum(Z_factor,2);

        if iter > INITTIME,

            for jj=1:J
                tt = GLOBAL_T(:,jj); % Y*V'
                ss = GLOBAL_S(:,jj); % V*V'
                
                zz = Z_tis(:,jj);

                % update slab : w/o (zz.*zz) due to different variational
                denom_u = ss(jj) + invtausq_u*GLOBAL_N(jj);

                if denom_u <= 1e-4, % avoid NaN
                    denom_u = 1e-4;
                end

                uu = tt - U*ss + slab_tis(:,jj) * ss(jj);
                uu = uu ./ denom_u;
                uu = uu + randn(Ntis,1,'single').*sqrt(sigsq./denom_u);

                U(:,jj) = zz .* uu;
                slab_tis(:,jj) = uu;

                % update Z
                % log_var_ratio = 0.5*(log(1+ss(jj)./(Ngene*invtausq_slab)) - log(sigsq));

                mass = -0.5 * uu.^2 * ss(jj) / sigsq;
                mass = mass + uu.*(tt - U*ss + (zz.*uu)*ss(jj)) /sigsq;
                mass = mass + prior_odds_tis;

                zz = 1./(1+exp(-mass));
                zz = single(rand(Ntis,1,'single') < zz);

                U(:,jj) = zz .* uu;
                Z_tis(:,jj) = zz;            
            end

        else
            for jj=1:J
                uu_old = U(:,jj);
                tt = GLOBAL_T(:,jj);
                ss = GLOBAL_S(:,jj);
                
                denom_u = ss(jj) + invtausq_u*GLOBAL_N(jj);

                if denom_u <= 1e-4, % avoid NaN
                    denom_u = 1e-4;
                end

                uu = (tt - U*ss + uu_old*ss(jj))/denom_u;
                uu = uu + randn(Ntis,1,'single') .* sqrt(sigsq/denom_u);
                
                U(:,jj) = uu;
            end
            slab_tis = U;
        end
        
        % ****************************************************************
        % sample prior odds for rank
        if iter > INITTIME,
            % r = betarnd(A0*Ngene+sum(Z_factor(:)),B0*Ngene+sum(1-Z_factor(:)));
            % prior_odds_factor = log(r) - log(1-r);

            r = betarnd(A0+sum(Z_tis,2),B0+sum(1-Z_tis,2));
            prior_odds_tis = log(r) - log(1-r);
        end

        sigsq = (GLOBAL_RSS + sigsq_prior*SIGSQ0*(J*Ntis + J*Ngene*Nind));
        sigsq = sigsq / (GLOBAL_OBS + J*Ntis + J*Ngene*Nind);
        
        mean_activity_tis = mean(Z_tis(:));
        mean_activity_rank = mean(Zt(:));

        % ================================================================
        if iter > BURNIN + INITTIME

            % model averaging
            Vavg(:,:,genes) = Vavg(:,:,genes) + reshape(Vt', J, Nind, ng) / (NGIBBS/Nbatch);
            VZavg(:,:,genes) = VZavg(:,:,genes) + reshape(VZt', J, Nind, ng) / (NGIBBS/Nbatch);
            intercept_avg(:,genes) = intercept_avg(:,genes) + intercept/(NGIBBS/Nbatch);
            rank_spike_avg(genes,:) = rank_spike_avg(genes,:) + Zt / (NGIBBS/Nbatch);

            if mod(iter,Nbatch) == 0,
                % rank_spike_samples(:,:,gibbs) = Z_factor';
                % invtausq_v_samples(genes,:,gibbs) = invtausq_v;
                u_spike_samples(:,:,gibbs) = Z_tis;
                u_slab_samples(:,:,gibbs) = slab_tis;
            end

            fprintf(2,'Gibbs = %03d, Iter = %05d, rmse = %.4e, sigma = %.4e, tis activity = %.4e, rank activity = %.4e\n', ...
                    epoch, iter, rmse, sqrt(sigsq), mean_activity_tis, mean_activity_rank);

        else

            fprintf(2,'Epoch = %03d, Iter = %05d, rmse = %.4e, sigma = %.4e, tis activity = %.4e, rank activity = %.4e\n', ...
                    epoch, iter, rmse, sqrt(sigsq), mean_activity_tis, mean_activity_rank);

        end

    end



    % ****************************************************************
    % final estimate of Vtot
    for bb = 0:(Nbatch-1),

        outfile = sprintf('%s/V_minibatch_%d.mat', OUTDIR, bb);
        genes_minibatch = GENES_minibatch{bb+1};
        V_minibatch = Vavg(:,:,genes_minibatch);
        VZ_minibatch = VZavg(:,:,genes_minibatch);
        save(outfile,'V_minibatch','VZ_minibatch','genes_minibatch','-v7.3');

    end

    U = mean(u_spike_samples.*u_slab_samples,3);

    save([OUTDIR,'/parameters.mat'], 'U', 'intercept_avg', 'rank_spike_avg', '*_samples', 'invtausq_u', 'prior_*', 'sigsq','rmse_trace', 'Ngene', 'Ntis', 'Nind', '-v7.3');

end
